from pylab import *  # The plotting module
import random
import copy

# ======================= #

generations500 = list()

# ======================= #

def is_allele_a_over(lst):
    if 'Aa' in lst:
        return False
    if 'aA' in lst:
        return False
    if 'aa' in lst:
        return False
    return True

# ======================= #

# a function to calculatge the count of alleles
# return count value: heterozygous, dominant, recessive
# heterozygous = Aa, aA; dominant = AA; recessive = aa;
def countGen_hdr(gen):
    heterozygous_cnt = 0
    dominant_cnt = 0
    recessive_cnt = 0

    for i in gen:
        if i == 'Aa' or i == 'aA':
            heterozygous_cnt += 1
        elif i == 'AA':
            dominant_cnt += 1
        elif i == 'aa':
            recessive_cnt += 1

    return (heterozygous_cnt, dominant_cnt, recessive_cnt)

# ======================= #

def func_run_500_gen():
    popu = list()

    popu = ["AA"] * 25
    popu.extend(["Aa"]*25)
    popu.extend(["aA"]*25)
    popu.extend(["aa"]*25)

    # ========== #

    generations500.append(copy.deepcopy(popu)) # initial generation

    # ========== #

    for generation in range(500):    # 500 generations
        individuals100 = list()
        individual = 0

        while individual < 100: # do until 100 individuals made

            pick1 = random.randint(0, 99)
            # do until pick a different number with pick1
            while True:
                pick2 = random.randint(0, 99)
                if pick1 != pick2:
                    break

            # 2 different individual number in range [0, 99] picked
            # combine to create a new individual
            combine = str()
            combine = generations500[generation][pick1][random.randint(0,1)]
            combine += (generations500[generation][pick2][random.randint(0,1)])

            if combine == "aa":
                # only 80% of ‘aa’ individuals will survive
                if random.random() > 0.8: # you can change 0.8 to 0, 0.2, 1.0 etc. to see what happen. then you understand it
                    print("sorry, this 'aa' is dead")
                    continue    # throw away this 'aa' combination. it's dead

            # good. all individuals (new_borns) are alive
            individuals100.append(copy.deepcopy(combine))

            individual += 1

        generations500.append(copy.deepcopy(individuals100)) # must use deepcopy

        if is_allele_a_over(individuals100): # use shallowcopy is enough
            print("\n\nallele 'a' disappears from population after", generation, "generation(s).\n\n")
            break

# ========================== main  ========================== #
func_run_500_gen()
heterozygous = list()
recessive = list()
dominant = list()

# count allele
for i in generations500:
    h, d, r = countGen_hdr(i)
    heterozygous.append(h)
    dominant.append(d)
    recessive.append(r)

print("generations500 detail {")
print("len = ", len(generations500))
for i, e in enumerate(generations500):
    print(i)
    print(e)
    print("--->")
print("generations500 detail }")

# ======= to plot ======= #
x_data = list()
for i in range(len(generations500)):
    x_data.append(i)

plot(x_data, heterozygous, 'r', label="hetero")    # Plot the data. The 'r' sets the colour to red
plot(x_data, dominant, 'g', label="dominant")
plot(x_data, recessive, 'b', label="recessive")

xlabel("Generations")                   # Add x axis labels
ylabel("Count")                         # Add y axis labels
title("Trend for allele distribution")  # Add a title
legend(loc='best')                      # Add the key

show()  # Display the plot
